import struct
from enum import Enum
import itertools

import argparse

'''
Big thanks to the following for the format reference:
https://github.com/libyal/libmdmp/blob/main/documentation/Minidump%20(MDMP)%20format.asciidoc#57-32-bit-memory-range-descriptor
https://github.com/skelsec/minidump
'''


classes_interest = ['fxDTSBrick', 'AIPlayer', 'Player', 'WheeledVehicle', 'FlyingVehicle', 'HoverVehicle']



def num_to_hex(n, leading_zeros):
	return "{0:#0{1}x}".format(n, leading_zeros + 2)
parser = argparse.ArgumentParser()

parser.add_argument('analyze_file', nargs='?', type=str, help = "analyze an MDMP file")
#parser.add_argument("-o", "--output-file", type=str, default=None, help = "output file for analysis results")

parser.add_argument("--analyze-objects", action='store_true', help = "Analyze objects from object dictionary; the default offset works for r2033")
parser.add_argument("--print-tracebuddy", action='store_true', help = "Print TraceBuddy journal; also define --tbd or --tbo")
#parser.add_argument("--print-objecttrace", action='store_true', help = "Print ObjectTrace journal; also define --otd or --oto")

parser.add_argument("--bl-object-offset", "--bloo", type=str, default='0x3C7234', help = "offset of object ID dictionary (tsf_gIdDictionaryLoc)")


parser.add_argument("--tracebuddy-direct-address", "--tbd", type=str, default=None, help = "direct pointer to trace_lines std::list; now recorded in console.log from TraceBuddy version 2.1")
parser.add_argument("--tracebuddy-offset", "--tbo", type=str, default=None, help = "offset of trace_lines std::list from image base of DLL")

#parser.add_argument("--objecttrace-direct-address", "--otd", type=str, default=None, help = "direct pointer to object_journal std::list; you are told this in console.log when ObjectTrace.dll initializes")
#parser.add_argument("--objecttrace-offset", "--oto", type=str, default=None, help = "offset of object_journal std::list from image base of DLL")


parser.add_argument("--help-offsets", action='store_true', help = "List common offsets")
args = parser.parse_args()

if args.help_offsets:
#	print('ObjectTrace.dll (all versions) print the direct address near top of console.log so you can use --otd')
	#print()
	print('Note: TraceBuddy versions 2.1.0 and above print the direct address near top of console.log so you can use --tbd')
	print('TraceBuddy.dll version 2.1.0')
	print('    md5sum: 990c6bc179524f2ba88c50bcd662fb62')
	print('    released: November 6, 2023')
	print('    offset: 0x100034')
	print()
	print('TraceBuddy.dll version 2.0.0')
	print('    md5sum: ed71583537dc6077f97500e897aa0ed9')
	print('    released: October 28, 2023')
	print('    offset: 0x100034')
	print()
	print('TraceBuddy.dll (Buddy\'s private version)')
	print('    md5sum: 2b41150e2a549d2053ad5db432f3b79f')
	print('    offset: 0xFB034')
	print()
	print('TraceBuddy.dll version 1.0.1')
	print('    md5sum: 03885bed372f753795b77e9b40da930d')
	print('    released: March 6, 2022')
	print('    offset: 0xC400C')
	print()
	print('TraceBuddy.dll version 1.0.0 did not use std::list.')
	exit(0)


args.bl_object_offset = int(args.bl_object_offset.replace("0x", ""), 16)


stdlist_routines = {\
	'TraceBuddy': {},
	#'ObjectTrace': {},
	}


stdlist_routines['TraceBuddy']['offset'] = args.tracebuddy_offset
stdlist_routines['TraceBuddy']['direct_address'] = args.tracebuddy_direct_address
stdlist_routines['TraceBuddy']['ptr_name'] = 'trace_lines'

'''
stdlist_routines['ObjectTrace']['offset'] = args.objecttrace_offset
stdlist_routines['ObjectTrace']['direct_address'] = args.objecttrace_direct_address
stdlist_routines['ObjectTrace']['ptr_name'] = 'journal_lines'
'''

for name, data in stdlist_routines.items():
	data['mode_direct_address'] = False
	data['mode_offset'] = False
	data['base_addr'] = 0
	data['detected_offset'] = 0

	#Neither - assume offset for TraceBuddy version 2.1.0
	if data['offset'] is None and data['direct_address'] is None:
		if name == 'TraceBuddy':
			data['offset'] = 0x100034
			data['mode_offset'] = True
		'''
		elif name == 'ObjectTrace':
			data['offset'] = 0x10105C
			data['mode_offset'] = True
		'''
	else:
		if data['offset'] is not None:
			data['offset'] = int(data['offset'].replace("0x", ""), 16)
			data['mode_offset'] = True

		if data['direct_address'] is not None:
			data['direct_address'] = int(data['direct_address'].replace("0x", ""), 16)
			data['mode_direct_address'] = True
			
	if data['mode_offset'] and data['mode_direct_address']:
		print('Can not use both offset and direct-address for', name)
		exit(1)

class MINIDUMP_STREAM_TYPE(Enum):
	UnusedStream = 0
	ThreadListStream = 3
	ModuleListStream = 4
	MemoryListStream = 5
	SystemInfoStream = 7
	Memory64ListStream = 9 #64-bit memory allocation information stream
	HandleDataStream = 12
	MiscInfoStream = 15
	MemoryInfoListStream = 16 #Memory region description information stream
	ThreadInfoListStream = 17
	Unknown21 = 21
	Unknown22 = 22

f = None
memory_range_list = []
memory_rva = 0


base_addr_blockland = 0

def seek_to_memory_ranged(addr):
	base_offset_in_file = memory_rva
	for memory_range in memory_range_list:
		if addr >= memory_range['rva'] and addr < (memory_range['rva'] + memory_range['size']):
			offset_in_range = addr - memory_range['rva']
			f.seek(base_offset_in_file + offset_in_range)
			return True
		base_offset_in_file += memory_range['size']

	return False

def read_stdlist_common(routine_name):
	routine_data = stdlist_routines[routine_name]

	list_trace_lines_ptr_loc = 0
	if routine_data['mode_offset']:
		list_trace_lines_ptr_loc = routine_data['offset'] + routine_data['base_addr']
	else:
		list_trace_lines_ptr_loc = routine_data['direct_address']
	print("Reading pointer at", num_to_hex(list_trace_lines_ptr_loc, 8))
	seek_to_memory_ranged(list_trace_lines_ptr_loc)
	
	if not list_trace_lines_ptr_loc:
		print('Could not locate memory block containing pointer')
		return
		
	if list_trace_lines_ptr_loc == 0:
		print('Null pointer')
		return
		
	list_trace_lines_ptr = struct.unpack('I', f.read(4))[0]
	print(routine_data['ptr_name'], "std::list starts at", num_to_hex(list_trace_lines_ptr, 8))
	addr_next = list_trace_lines_ptr
	
	#Read first element
	seek_to_memory_ranged(addr_next + 4)
	addr_prev = struct.unpack('I', f.read(4))[0]
	
	if addr_prev == list_trace_lines_ptr_loc:
		print('Located first element')
		print()

	line_count=0
	while addr_next != list_trace_lines_ptr_loc:
		seek_to_memory_ranged(addr_next)
		addr_next = struct.unpack('I', f.read(4))[0]
		addr_prev = struct.unpack('I', f.read(4))[0]
		addr_str = struct.unpack('I', f.read(4))[0]
		str_len = struct.unpack('I', f.read(4))[0]

		seek_to_memory_ranged(addr_str)

		print(f.read(str_len).decode('cp1252'), end='')
		line_count += 1
		
	print()
	print('Found', line_count, 'lines of data for', routine_data['ptr_name'], 'in', routine_name)
	if routine_data['mode_direct_address']:
		calc_offset = hex(routine_data['detected_offset'])[2:].upper()
		print('For dumps with this same DLL, use this offset instead of reading the pointer from the console log: 0x'+ str(calc_offset))

def read_objects():
	all_object_addresses = {}

	tsf_gIdDictionaryPtrLoc = base_addr_blockland + args.bl_object_offset
	seek_to_memory_ranged(tsf_gIdDictionaryPtrLoc)
	tsf_gIdDictionaryPtr = struct.unpack('I', f.read(4))[0]

	for bucket_num in range(0, 4096):
		seek_to_memory_ranged(tsf_gIdDictionaryPtr + bucket_num * 4)
		next_obj_ptr = struct.unpack('I', f.read(4))[0]

		while True:
			if next_obj_ptr == 0:
				break
			seek_to_memory_ranged(next_obj_ptr)

			obj_ptr = next_obj_ptr
			f.read(16)
		
			#SimObject* SimObject->nextIdObject
			next_obj_ptr = struct.unpack('I', f.read(4))[0]
			f.read(12)
			obj_id = struct.unpack('i', f.read(4))[0]

			all_object_addresses[obj_id] = obj_ptr
			

	for obj_id in sorted(all_object_addresses.keys()):
		obj_addr = all_object_addresses[obj_id]
		seek_to_memory_ranged(obj_addr)
		vt=struct.unpack('I', f.read(4))[0]
		obj_name_ptr=struct.unpack('I', f.read(4))[0]
		
		f.read(8)
		
		#SimObject* SimObject->nextIdObject
		next_obj_ptr = struct.unpack('I', f.read(4))[0]

		f.read(12)
		obj_id = struct.unpack('i', f.read(4))[0]
		namespace_addr = struct.unpack('I', f.read(4))[0]
		
		print('Object', num_to_hex(obj_addr, 8))
		
		obj_name=b''
		if obj_name_ptr != 0:
			seek_to_memory_ranged(obj_name_ptr)
			next_char = f.read(1)
			while next_char != b'\0':
				obj_name += next_char
				next_char = f.read(1)
			obj_name = obj_name.decode('utf-8')
			print('      Name:', obj_name)
		print('        id:', obj_id)
		ns_name='NULL'
		this_namespace_addr = namespace_addr
		while this_namespace_addr != 0:
			seek_to_memory_ranged(this_namespace_addr)
			
			ptr_ns_name=struct.unpack('I', f.read(4))[0]
			f.read(4)
			ptr_mParent = struct.unpack('I', f.read(4))[0]
			f.read(4)
			ptr_mClassRep = struct.unpack('I', f.read(4))[0]
			
			if ptr_mClassRep != 0:
				if ptr_ns_name == 0:
					break
				seek_to_memory_ranged(ptr_ns_name)
				next_char = f.read(1)
				ns_name=b''
				while next_char != b'\0':
					ns_name += next_char
					next_char = f.read(1)
				ns_name = ns_name.decode('utf-8')
				break
			this_namespace_addr = ptr_mParent
		print('     class:', ns_name)
		
		if ns_name in classes_interest:
			seek_to_memory_ranged(obj_addr + 92)
			#obj->mObjToWorld
			
			matrix_f_vals = []
			for i in range(0, 16):
				matrix_f_vals += [str(struct.unpack('f', f.read(4))[0])]
				
			pos = ' '.join([matrix_f_vals[3], matrix_f_vals[7], matrix_f_vals[11]])
			print('       pos:', pos)
		print()
			
		
	
with open(args.analyze_file, "rb") as f:

	print_stack_locs = []

	header_sig = f.read(4)
	if header_sig != b'MDMP':
		print("Not an MDMP file")
		exit(1)
	version = f.read(2)
	version_imp = f.read(2)
	num_streams = struct.unpack('I', f.read(4))[0]
	stream_dir_RVA = struct.unpack('I', f.read(4))[0]
	checksum = f.read(4)
	timestamp = f.read(4)
	flags = struct.unpack('Q', f.read(8))[0]
	
	#Flags set:
	#MiniDumpWithFullMemory
	#MiniDumpWithHandleData
	#MiniDumpWithFullMemoryInfo

	for stream_num in range(0, num_streams):
		f.seek(stream_dir_RVA + stream_num * 4 * 3, 0 )
		stream_type = struct.unpack('I', f.read(4))[0]
		
		stream_data_size = struct.unpack('I', f.read(4))[0]
		stream_RVA = struct.unpack('I', f.read(4))[0]
	
		if stream_type not in \
			[MINIDUMP_STREAM_TYPE.ModuleListStream.value,\
			 MINIDUMP_STREAM_TYPE.ThreadListStream.value,\
			 MINIDUMP_STREAM_TYPE.ThreadInfoListStream.value,\
			 MINIDUMP_STREAM_TYPE.Memory64ListStream.value]:
			continue
			
		print(MINIDUMP_STREAM_TYPE(stream_type).name, 'size =', stream_data_size, 'location =', stream_RVA)

		f.seek(stream_RVA, 0)
		
		if stream_type == MINIDUMP_STREAM_TYPE.ThreadListStream.value:
			num_threads = struct.unpack('i', f.read(4))[0]
			print('    Num threads:', num_threads)
			tell_pos = f.tell()
			for thread_index_num in range(0, num_threads):
				f.seek(tell_pos)
				
				thread_id = struct.unpack('i', f.read(4))[0]
				suspend_count = struct.unpack('i', f.read(4))[0]
				priority_class = struct.unpack('i', f.read(4))[0]
				priority = struct.unpack('i', f.read(4))[0]
				teb = struct.unpack('Q', f.read(8))[0]
				stack_data_size = struct.unpack('i', f.read(4))[0]
				stack_rva = struct.unpack('i', f.read(4))[0]
				thread_context_data_size = struct.unpack('i', f.read(4))[0]
				thread_context_rva = struct.unpack('i', f.read(4))[0]
					
				tell_pos = f.tell()
				
				if thread_context_rva != 0:
					f.seek(thread_context_rva)
					ctx = {}
					print(thread_id, num_to_hex(teb, 16), num_to_hex(stack_rva, 8), num_to_hex(thread_context_rva, 8))
					#print(f.read(thread_context_data_size))
					
					'''
					self.file_handle.seek(rva)
					if self.sysinfo.ProcessorArchitecture == PROCESSOR_ARCHITECTURE.AMD64:
						thread.ContextObject = CONTEXT.parse(self.file_handle)
					elif self.sysinfo.ProcessorArchitecture == PROCESSOR_ARCHITECTURE.INTEL:
						thread.ContextObject = WOW64_CONTEXT.parse(self.file_handle)
					'''
					ctx['P1Home'] = struct.unpack('Q', f.read(8))[0]
					ctx['P2Home'] = struct.unpack('Q', f.read(8))[0]
					ctx['P3Home'] = struct.unpack('Q', f.read(8))[0]
					ctx['P4Home'] = struct.unpack('Q', f.read(8))[0]
					ctx['P5Home'] = struct.unpack('Q', f.read(8))[0]
					ctx['P6Home'] = struct.unpack('Q', f.read(8))[0]
					
					ctx['ContextFlags'] = struct.unpack('i', f.read(4))[0]
					ctx['MxCsr'] = struct.unpack('i', f.read(4))[0]
					ctx['SegCs'] = struct.unpack('H', f.read(2))[0]
					ctx['SegDs'] = struct.unpack('H', f.read(2))[0]
					ctx['SegEs'] = struct.unpack('H', f.read(2))[0]
					ctx['SegFs'] = struct.unpack('H', f.read(2))[0]
					ctx['SegGs'] = struct.unpack('H', f.read(2))[0]
					ctx['SegSs'] = struct.unpack('H', f.read(2))[0]
					ctx['EFlags'] = struct.unpack('i', f.read(4))[0]
					ctx['Dr0'] = struct.unpack('Q', f.read(8))[0]
					ctx['Dr1'] = struct.unpack('Q', f.read(8))[0]
					ctx['Dr2'] = struct.unpack('Q', f.read(8))[0]
					ctx['Dr3'] = struct.unpack('Q', f.read(8))[0]
					ctx['Dr6'] = struct.unpack('Q', f.read(8))[0]
					ctx['Dr7'] = struct.unpack('Q', f.read(8))[0]
					ctx['Rax'] = struct.unpack('Q', f.read(8))[0]
					ctx['Rcx'] = struct.unpack('Q', f.read(8))[0]
					ctx['Rdx'] = struct.unpack('Q', f.read(8))[0]
					ctx['Rbx'] = struct.unpack('Q', f.read(8))[0]
					ctx['Rsp'] = struct.unpack('Q', f.read(8))[0]
					ctx['Rbp'] = struct.unpack('Q', f.read(8))[0]
					ctx['Rsi'] = struct.unpack('Q', f.read(8))[0]
					ctx['Rdi'] = struct.unpack('Q', f.read(8))[0]
					ctx['R8'] = struct.unpack('Q', f.read(8))[0]
					ctx['R9'] = struct.unpack('Q', f.read(8))[0]
					ctx['R10'] = struct.unpack('Q', f.read(8))[0]
					ctx['R11'] = struct.unpack('Q', f.read(8))[0]
					ctx['R12'] = struct.unpack('Q', f.read(8))[0]
					ctx['R13'] = struct.unpack('Q', f.read(8))[0]
					ctx['R14'] = struct.unpack('Q', f.read(8))[0]
					ctx['R15'] = struct.unpack('Q', f.read(8))[0]
					ctx['Rip'] = struct.unpack('Q', f.read(8))[0]
					print(ctx.items())
					print('rip:', num_to_hex(ctx['Rip'], 16))
					print('rsp:', num_to_hex(ctx['Rsp'], 16))
					if ctx['Rsp'] != 0:
						print_stack_locs += [ctx['Rsp']]
					'''

					ctx.DUMMYUNIONNAME = CTX_DUMMYUNIONNAME.parse(buff)
					
					ctx.VectorRegister = M128A.parse_array(buff, 26)         # M128A   [26]
					ctx.VectorControl =  int.from_bytes(buff.read(8), byteorder = 'little', signed = False)       # DWORD64
					ctx.DebugControl = int.from_bytes(buff.read(8), byteorder = 'little', signed = False)         # DWORD64
					ctx.LastBranchToRip = int.from_bytes(buff.read(8), byteorder = 'little', signed = False)      # DWORD64
					ctx.LastBranchFromRip = int.from_bytes(buff.read(8), byteorder = 'little', signed = False)    # DWORD64
					ctx.LastExceptionToRip = int.from_bytes(buff.read(8), byteorder = 'little', signed = False)   # DWORD64
					ctx.LastExceptionFromRip = int.from_bytes(buff.read(8), byteorder = 'little', signed = False) # DWORD64
					'''
		'''
		if stream_type == MINIDUMP_STREAM_TYPE.ThreadInfoListStream.value:
			size_of_header = struct.unpack('i', f.read(4))[0]
			size_of_entry = struct.unpack('i', f.read(4))[0]
			num_entries = struct.unpack('i', f.read(4))[0]
			for thread_entry_num in range(0, num_entries):
				thread_id = struct.unpack('i', f.read(4))[0]
				dump_flags = struct.unpack('i', f.read(4))[0]
				dump_error = struct.unpack('i', f.read(4))[0]
				exit_status = struct.unpack('i', f.read(4))[0]
				create_time = struct.unpack('Q', f.read(8))[0]
				exit_time = struct.unpack('Q', f.read(8))[0]
				kernel_time = struct.unpack('Q', f.read(8))[0]
				user_time = struct.unpack('Q', f.read(8))[0]
				start_address = struct.unpack('Q', f.read(8))[0]
				affinity = struct.unpack('Q', f.read(8))[0]
		'''
		
		if stream_type == MINIDUMP_STREAM_TYPE.ModuleListStream.value:
			num_modules = struct.unpack('I', f.read(4))[0]
			print('    Num modules:', num_modules)
			
			module_list_base = f.tell()
			for module_num in range(0, num_modules):
				f.seek(module_list_base + module_num * 108, 0)
				image_base = struct.unpack('Q', f.read(8))[0]
				image_size = struct.unpack('I', f.read(4))[0]
				module_checksum = struct.unpack('I', f.read(4))[0]
				module_timestamp = struct.unpack('I', f.read(4))[0]
				module_name_rva = struct.unpack('I', f.read(4))[0]
				module_version_info = f.read(52)
				
				f.seek(module_name_rva)
				module_name_size = struct.unpack('I', f.read(4))[0]
				module_name_str_full = f.read(module_name_size).decode('utf-16-le')
				module_name_str = module_name_str_full.split('\\')[-1]
				print('        ', num_to_hex(image_base, 8), image_size, module_name_size, module_name_str_full)
				if module_name_str == 'TraceBuddy.dll': # or module_name_str == 'ObjectTrace.dll':
					routine_name = module_name_str[0:-4]
					routine = stdlist_routines[routine_name]
					print('            Found '+ module_name_str)
					routine['base_addr'] = image_base
					if routine['mode_direct_address']:
						detected_offset = routine['direct_address'] - image_base
						routine['detected_offset'] = detected_offset
						calc_offset = hex(detected_offset)[2:].upper()
						print('            Back-calculated the offset: 0x'+ str(calc_offset))
				if module_name_str == 'Blockland.exe':
					print('            Found Blockland.exe')
					base_addr_blockland = image_base
						

		if stream_type == MINIDUMP_STREAM_TYPE.Memory64ListStream.value:
			m64list_num_ranges = struct.unpack('Q', f.read(8))[0]
			m64list_rva = struct.unpack('Q', f.read(8))[0]
			memory_rva = m64list_rva
			memory_range_total_sizes = 0
			for memory_range_num in range(0, m64list_num_ranges):
				memory_range_rva = struct.unpack('Q', f.read(8))[0]
				memory_range_size = struct.unpack('Q', f.read(8))[0]
				memory_range_total_sizes += memory_range_size
				memory_range_list += [{'rva': memory_range_rva, 'size': memory_range_size}]
			print('    Found', memory_range_total_sizes, 'bytes of memory in', m64list_num_ranges, 'sections')

	print()
	
	for stack_loc in print_stack_locs:
		print(num_to_hex(stack_loc, 8))

		for i in range(0, 16):
			for j in range(0, 16):
				seek_to_memory_ranged(stack_loc - (i * 16 + j) * 4)
				print(num_to_hex(struct.unpack('I', f.read(4))[0], 1).rjust(10), end=' ')
			print()
	
	if args.print_tracebuddy:
		read_stdlist_common('TraceBuddy')
	
	#if args.print_objecttrace:
	#	read_stdlist_common('ObjectTrace')

	if args.analyze_objects:
		read_objects()
	
